#!/usr/bin/env python3
"""
YOLOv8 OpenVINO Go Service Client

Example client for interacting with the inference service.
"""

import requests
import json
import argparse
import sys
from pathlib import Path


class YOLOv8Client:
    def __init__(self, base_url="http://localhost:8080"):
        self.base_url = base_url
        self.session = requests.Session()

    def health_check(self):
        """Check service health"""
        try:
            response = self.session.get(f"{self.base_url}/health")
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Health check failed: {e}")
            return None

    def infer(self, image_path):
        """Run inference on a single image"""
        if not Path(image_path).exists():
            print(f"Error: Image not found: {image_path}")
            return None

        payload = {"image_path": image_path}
        try:
            response = self.session.post(
                f"{self.base_url}/api/v1/infer",
                json=payload
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Inference failed: {e}")
            return None

    def infer_batch(self, image_paths):
        """Run inference on multiple images"""
        # Verify all images exist
        for path in image_paths:
            if not Path(path).exists():
                print(f"Error: Image not found: {path}")
                return None

        payload = {"image_paths": image_paths}
        try:
            response = self.session.post(
                f"{self.base_url}/api/v1/infer/batch",
                json=payload
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Batch inference failed: {e}")
            return None

    def get_model_info(self):
        """Get model information"""
        try:
            response = self.session.get(f"{self.base_url}/api/v1/model/info")
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Failed to get model info: {e}")
            return None

    def print_results(self, result):
        """Pretty print inference results"""
        if result is None:
            return

        print(json.dumps(result, indent=2))


def main():
    parser = argparse.ArgumentParser(
        description="YOLOv8 OpenVINO Go Service Client"
    )
    parser.add_argument(
        "--url",
        default="http://localhost:8080",
        help="Service URL (default: http://localhost:8080)"
    )
    parser.add_argument(
        "--health",
        action="store_true",
        help="Check service health"
    )
    parser.add_argument(
        "--info",
        action="store_true",
        help="Get model information"
    )
    parser.add_argument(
        "--infer",
        type=str,
        help="Run inference on image"
    )
    parser.add_argument(
        "--batch",
        nargs="+",
        help="Run batch inference on multiple images"
    )

    args = parser.parse_args()

    client = YOLOv8Client(args.url)

    if args.health:
        print("Checking service health...")
        result = client.health_check()
        client.print_results(result)
        return 0 if result else 1

    if args.info:
        print("Getting model information...")
        result = client.get_model_info()
        client.print_results(result)
        return 0 if result else 1

    if args.infer:
        print(f"Running inference on {args.infer}...")
        result = client.infer(args.infer)
        client.print_results(result)
        if result and result.get("status") == "success":
            print(f"\nDetections: {len(result.get('detections', []))}")
            for det in result.get("detections", []):
                print(f"  - {det['class_name']}: {det['confidence']:.2%}")
        return 0 if result else 1

    if args.batch:
        print(f"Running batch inference on {len(args.batch)} images...")
        result = client.infer_batch(args.batch)
        client.print_results(result)
        return 0 if result else 1

    parser.print_help()
    return 0


if __name__ == "__main__":
    sys.exit(main())

